
#include "heaviside_tiling.h"
#include "register/op_def_registry.h"

#define BLOCK_BYTES_SIZE 32
#define CORENUM 40

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context) {
    HeavisideTilingData tiling;
    const gert::StorageShape* input_shape = context->GetInputShape(0);
    uint32_t inputSize = 1;
    for (int i = 0; i < input_shape->GetStorageShape().GetDimNum(); i++) {
        inputSize *= input_shape->GetStorageShape().GetDim(i);
    }
        
    const gert::StorageShape* value_shape = context->GetInputShape(1);
    uint32_t valueSize = 1;
    for (int i = 0; i < value_shape->GetStorageShape().GetDimNum(); i++) {
        valueSize *= value_shape->GetStorageShape().GetDim(i);
    }

    // printf("inputSize: %d, valueSize: %d\n",inputSize,valueSize);
    tiling.set_inputSize(inputSize);
    tiling.set_valueSize(valueSize);

    uint32_t inputBytes = GetSizeByDataType(context->GetInputDesc(0)->GetDataType());

    if (valueSize == 1)
    {
        int BlockSize = BLOCK_BYTES_SIZE / inputBytes;
        int BlockNum = (inputSize + BlockSize - 1)/ BlockSize;
        int coreNum = BlockNum < CORENUM ? BlockNum : CORENUM ;
        // printf("BlockNum=%d, coreNum=%d\n",BlockNum,coreNum);
        context->SetBlockDim(coreNum);
        context->SetTilingKey(1);
    } else {
        context->SetBlockDim(1);
        context->SetTilingKey(2);
    }

    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
                        context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    return ge::GRAPH_SUCCESS;
}
}  // namespace optiling

// namespace ge {
// static ge::graphStatus InferShape(gert::InferShapeContext* context)
// {
//     const gert::Shape* x1_shape = context->GetInputShape(0);
//     gert::Shape* y_shape = context->GetOutputShape(0);
//     *y_shape = *x1_shape;
//     return GRAPH_SUCCESS;
// }
// }

namespace ops {
class Heaviside : public OpDef {
public:
    explicit Heaviside(const char* name) : OpDef(name)
    {
        this->Input("input")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("values")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("out")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        // this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(Heaviside);
}
